import numpy as np
import pandas as pd
import re
import json
import logging
import os

import admm.ADMM_Pregel as ADMM
from admm.pregel import Pregel
import admm.ADMM_update as ADMM_update

import graphscope
from graphscope.framework.loader import Loader

"""
阻抗单位 欧姆
电流单位 kA
电压单位 kV
容量/有功/无功单位 Mva/MW/Mvar
能量 MWh
价格系数 元/MW
"""


class acline_args_pd:
    args: dict[str, dict]

    def __init__(self, line_table: pd.DataFrame, bus_table: pd.DataFrame) -> None:
        b1 = bus_table.rename(
            columns={
                "name": "bus1",
                "kV": "kV_bus1",
                "vminpu": "vminpu_bus1",
                "vmaxpu": "vmaxpu_bus1",
            },
            copy=False,
        )
        df = pd.merge(line_table, b1, how="inner", on="bus1", copy=False)
        b2 = bus_table.rename(
            columns={
                "name": "bus2",
                "kV": "kV_bus2",
                "vminpu": "vminpu_bus2",
                "vmaxpu": "vmaxpu_bus2",
            },
            copy=False,
        )
        df = pd.merge(df, b2, how="inner", on="bus2", copy=True)
        df.set_index("name", inplace=True)
        df["Vmax_i"] = df["kV_bus1"] * df["vmaxpu_bus1"]
        df["Vmin_i"] = df["kV_bus1"] * df["vminpu_bus1"]
        df["Vmax_j"] = df["kV_bus2"] * df["vmaxpu_bus2"]
        df["Vmin_j"] = df["kV_bus2"] * df["vminpu_bus2"]
        self.args = df.to_dict(orient="index")

    @staticmethod
    def allocate_initialize_value(
        acline_args: "acline_args_pd", T: int
    ) -> dict[str, dict[str,ADMM.terminal_args]]:
        init_args = [
            (
                key,
                {
                    "init_terminal1": ADMM.terminal_args(T, ("P", "Q", "V")),
                    "init_terminal2": ADMM.terminal_args(T, ("P", "Q", "V")),
                },
            )
            for key in acline_args.args.keys()
        ]
        init_args = dict(init_args)
        return init_args

    @staticmethod
    def generate_zeros_initialize(
        acline_args: "acline_args_pd", T: int
    ) -> dict[str, dict[str, ADMM.terminal_args]]:
        # 将P、Q及其拉格朗日乘子设为0，将V设为对应额定电压，V的拉格朗日乘子为0
        # 返回值可以给set_initialize_value
        init_args = acline_args_pd.allocate_initialize_value(acline_args, T)
        for k, v in init_args.items():
            v["init_terminal1"].P = np.zeros((T, 1), dtype=np.float64)
            v["init_terminal2"].P = np.zeros((T, 1), dtype=np.float64)
            v["init_terminal1"].Q = np.zeros((T, 1), dtype=np.float64)
            v["init_terminal2"].Q = np.zeros((T, 1), dtype=np.float64)
            v["init_terminal1"].V2 = (
                np.ones((T, 1), np.float64) * acline_args.args[k]["kV_bus1"] ** 2
            )
            v["init_terminal2"].V2 = (
                np.ones((T, 1), np.float64) * acline_args.args[k]["kV_bus2"] ** 2
            )
            v["init_terminal1"].LP = np.zeros((T, 1), dtype=np.float64)
            v["init_terminal2"].LP = np.zeros((T, 1), dtype=np.float64)
            v["init_terminal1"].LQ = np.zeros((T, 1), dtype=np.float64)
            v["init_terminal2"].LQ = np.zeros((T, 1), dtype=np.float64)
            v["init_terminal1"].LV2 = np.zeros((T, 1), dtype=np.float64)
            v["init_terminal2"].LV2 = np.zeros((T, 1), dtype=np.float64)
        return init_args

    def set_initialize_value(self, init_val: dict[str, dict[str, ADMM.terminal_args]]):
        # 给所有线路设置初始值
        for k, v in init_val.items():
            self.args[k].update(v)

    def query(self, name: str) -> dict:
        return self.args[name]


class acload_args_pd:
    args: dict[str, dict]

    def __init__(self, acload_tbl: pd.DataFrame) -> None:
        self.args = {}
        for k in acload_tbl.index:
            Pmulti = acload_tbl.loc[k, "Pmultistr"].strip()
            Pmulti = re.sub(r"[\n\r]", r"", Pmulti)
            Pmulti = re.sub(r" {2,}", r" ", Pmulti)
            Pmulti = Pmulti.split(" ")
            Pmulti = np.array(Pmulti, dtype=np.float64)
            Pt = acload_tbl.loc[k, "ratedP"] * Pmulti

            Qmulti = acload_tbl.loc[k, "Qmultistr"].strip()
            Qmulti = re.sub(r"[\n\r]", r"", Qmulti)
            Qmulti = re.sub(r" {2,}", r" ", Qmulti)
            Qmulti = Qmulti.split(" ")
            Qmulti = np.array(Qmulti, dtype=np.float64)
            Qt = acload_tbl.loc[k, "ratedP"] * Qmulti

            self.args.update({acload_tbl.loc[k, "name"]: {"Pt": Pt, "Qt": Qt}})

    def query(self, name: str) -> dict:
        return self.args[name]


class storage_args_pd:
    args: dict[str, dict]

    def __init__(self, sto_tbl: pd.DataFrame) -> None:
        stbl=sto_tbl.set_index("name")
        self.args = stbl.to_dict(orient="index")

    def query(self, name: str) -> dict:
        return self.args[name]

    @staticmethod
    def allocate_initialize_value(
        storage_args: "storage_args_pd", T: int
    ) -> dict[str, ADMM.terminal_args]:
        init_args = [
            (
                key,
                ADMM.terminal_args(T, ("P", "Q")),
            )
            for key in storage_args.args.keys()
        ]
        init_args = dict(init_args)
        return init_args

    @staticmethod
    def generate_zeros_initialize(
        storage_args: "storage_args_pd", T: int
    ) -> dict[str, ADMM.terminal_args]:
        # 将P、Q及其拉格朗日乘子设为0
        # 返回值可以给set_initialize_value
        init_args = storage_args_pd.allocate_initialize_value(storage_args, T)
        for k, v in init_args.items():
            v.P = np.zeros((T, 1), dtype=np.float64)
            v.Q = np.zeros((T, 1), dtype=np.float64)
            v.LP = np.zeros((T, 1), dtype=np.float64)
            v.LQ = np.zeros((T, 1), dtype=np.float64)
        return init_args

    def set_initialize_value(self, init_val: dict[str, ADMM.terminal_args]):
        # 给所有储能设置初始值
        for k, v in init_val.items():
            self.args[k].update({"init_terminal":v})


class upper_args_pd:
    args: dict[str, dict]

    def __init__(self, src_tbl: pd.DataFrame, T: int) -> None:
        stbl=src_tbl.set_index("name")
        self.args = stbl.to_dict(orient="index")
        for k in self.args.keys():
            self.args[k]["Ct"] = (
                np.ones((T, 1), dtype=np.float64) * self.args[k]["price"]
            )

    def query(self, name: str) -> dict:
        return self.args[name]

    @staticmethod
    def allocate_initialize_value(
        upper_grid_property: "upper_args_pd", T: int
    ) -> dict[str, ADMM.terminal_args]:
        init_args = [
            (
                key,
                ADMM.terminal_args(T, ("P", "Q","V"),("V")),
            )
            for key in upper_grid_property.args.keys()
        ]
        init_args = dict(init_args)
        return init_args

    @staticmethod
    def generate_zeros_initialize(
        upper_grid_property: "upper_args_pd", T: int
    ) -> dict[str, ADMM.terminal_args]:
        # 将P、Q及其拉格朗日乘子设为0
        # 返回值可以给set_initialize_value
        init_property = upper_args_pd.allocate_initialize_value(upper_grid_property, T)
        for k, v in init_property.items():
            v.P = np.zeros((T, 1), dtype=np.float64)
            v.Q = np.zeros((T, 1), dtype=np.float64)
            v.LP = np.zeros((T, 1), dtype=np.float64)
            v.LQ = np.zeros((T, 1), dtype=np.float64)
            v.V2 = np.ones((T, 1), dtype=np.float64) * (upper_grid_property.args[k]["kV"] **2)
        return init_property

    def set_initialize_value(self, init_property: dict[str, ADMM.terminal_args]):
        # 给所有电网设置初始值
        for k, v in init_property.items():
            self.args[k].update({"init_terminal": v})


acline_tbl = pd.DataFrame(
    [
        ("L12", "b1", "b2", 0.3061 * 2000 / 5280.0, 0.627 * 2000 / 5280.0, None),
        ("L23", "b2", "b3", 0.01 * 12.47**2 / 6.0, 0.06 * 12.47**2 / 6.0, None),
        ("L34", "b3", "b4", 0.3061 * 2500 / 5280.0, 0.627 * 2500 / 5280.0, None),
    ],
    columns=["name", "bus1", "bus2", "R", "X", "Imax"],
)

bus_table = pd.DataFrame(
    [
        ("b1", 12.47, 1.10, 0.9),
        ("b2", 12.47, 1.10, 0.9),
        ("b3", 12.47, 1.10, 0.9),
        ("b4", 12.47, 1.10, 0.9),
    ],
    columns=["name", "kV", "vmaxpu", "vminpu"],
)

# 恒定购电成本
upper_table = pd.DataFrame(
    [("src", "b1", 12.47, 1800.0 * 3 * 10 / 1000, 0, 1.0 / 1000)],
    columns=["name", "bus1", "kV", "Pmax", "Pmin", "price"],
)
storage_tbl = pd.DataFrame(
    [("sto1", "b3", 0.5, 0.98, 0.98, 1.0, 0.99 / 1000, None, 4)],
    columns=[
        "name",
        "bus1",
        "Pmax",
        "eff_ch",
        "eff_dis",
        "capacity",
        "operating cost",
        "changing cost",
        "max changing",
    ],
)


Pstr = """ 
0.092 0.092 0.092 0.092 0.092 0.092 0.092 0.092 0.092 0.092 0.092 
0.174 0.174 0.174 0.174 0.174 2.23 2.23 2.23 0.174 0.174 0.174 
0.191 0.251 0.251 0.169 0.169 0.169 0.269 0.329 0.555 0.535 0.535 
0.57 0.57 0.57 0.57 0.57 0.57 0.57 0.57 0.57 0.488 0.488 0.743 
0.743 0.743 0.743 0.743 0.743 0.659 0.659 0.599 0.599 0.646 0.646 
0.68 0.68 0.74 0.74 0.381 0.381 0.381 0.381 0.381 0.381 0.215 0.215 
0.215 0.215 0.052 0.052 0.052 0.142 0.142 0.142 0.142 0.142 0.142 0.142 
0.142 0.142 0.142 0.142 0.142 0.142 0.142 0.142 0.142 0.142 0.142 0.052 
0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.196 0.196 0.196 
0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 
0.362 0.362 0.362 0.362 0.362 0.218 0.218 0.218 0.218 0.218 0.218 0.218 
0.218 0.218 0.218 0.308 0.471 0.471 0.471 0.471 0.471 0.471 0.305 0.305 
0.305 0.305 0.305 0.305 0.305 0.305 0.142 0.142 0.142 0.052 0.052 0.052 
0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 
0.052 0.052 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 
0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.052 0.052 
0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 
0.052 0.052 0.052 0.052 0.142 0.142 0.142 0.308 """

Qstr = """
0.343 0.343 0.343 0.447 0.447 0.447 0.507 0.507 0.507 0.507 0.507 0.507 
0.467 0.467 0.467 0.467 0.467 0.467 0.052 0.052 0.218 0.218 0.218 0.218 
0.218 0.218 0.218 0.218 0.218 0.218 0.218 0.218 0.218 0.218 0.218 0.218 
0.218 0.218 0.218 0.218 0.218 0.218 0.052 0.052 0.052 0.052 0.305 0.305 
0.305 0.305 0.305 0.305 0.305 0.305 0.305 0.305 0.305 0.305 0.305 0.305 
0.142 0.142 0.142 0.142 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 
0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 0.196 
0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 
0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.218 0.218 0.218 0.218 0.218 
0.218 0.218 0.218 0.218 0.218 0.218 0.218 0.218 0.381 0.381 0.381 0.381 
0.381 0.381 0.381 0.381 0.381 0.215 0.215 0.215 0.215 0.215 0.052 0.052 
0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.196 
0.196 0.196 0.196 0.196 0.196 0.286 0.286 0.286 0.286 0.286 0.286 0.286 
0.286 0.286 0.286 0.286 0.286 0.286 0.286 0.142 0.142 0.142 0.142 0.052 
0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 
0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.052 0.381 
0.381 0.381 0.381 0.381 0.381 0.381 0.614 0.614 0.614 0.696 0.716 1.029 
1.066 0.903 0.903 0.57 0.583 0.583 0.583 0.583 0.583 0.417 0.507 0.507 
0.507 0.651 0.711 0.711 0.711 0.711 0.711 0.711 0.711 0.711 0.711 0.711 
.711 0.711 0.711 0.711 0.621 0.654 0.654 0.654 0.754 0.61 0.406 0.406 
0.346 0.346 0.346 0.346 0.346 0.346 0.346 0.284 0.284 0.284 0.284 0.366 
0.366 0.366 0.284 0.324 0.324 0.324 0.324 0.324 0.304 0.304 0.304 0.304 
0.304 0.304 0.304 0.304 0.291 0.251 0.251 0.251 0.251 0.251 0.414 0.414 
0.65 0.633 0.633 0.633 0.633 0.633 0.715 0.715 1.048 1.048 1.048 1.192 
0.696 0.713 0.713 0.733 0.773 0.753 0.703 0.723 0.743 0.738 0.572 0.532 
"""

load_tbl = pd.DataFrame(
    [
        (
            "LD1",
            "b4",
            1800.0 / 1000 * 3,
            1800.0 / 1000 * 3 * np.sqrt((1 / 0.9) ** 2 - 1),
            Pstr,
            Qstr,
        )
    ],
    columns=["name", "bus1", "ratedP", "ratedQ", "Pmultistr", "Qmultistr"],
)


def append_id(tbl, vid_seed, num, dtype=np.int64):
    ids = pd.Series(
        np.arange(start=vid_seed, stop=vid_seed + num, step=1, dtype=dtype),
        name="vid",
    )
    tbl = pd.concat([ids, tbl], axis=1, copy=False)
    return tbl


def build_grid(
    acbus: pd.DataFrame,
    acline: pd.DataFrame,
    acload: pd.DataFrame,
    acsource: pd.DataFrame,
    storage: pd.DataFrame,
):

    sess = graphscope.session(cluster_type="hosts", num_workers=1, mode="eager")

    vid_seed = 10000
    acbus = append_id(acbus, vid_seed, acbus.shape[0], dtype=np.int64)
    vid_seed += acbus.shape[0]

    acline = append_id(acline, vid_seed, acline.shape[0], dtype=np.int64)
    vid_seed += acline.shape[0]

    acload = append_id(acload, vid_seed, acload.shape[0], dtype=np.int64)
    vid_seed += acload.shape[0]

    acsource = append_id(acsource, vid_seed, acsource.shape[0], dtype=np.int64)
    vid_seed += acsource.shape[0]

    storage = append_id(storage, vid_seed, storage.shape[0], dtype=np.int64)
    vid_seed += storage.shape[0]

    # oid是指acbus acline这些表里面id的数据类型，vid是graphscope内部存储的类型
    # 马的不支持空值，只要输入的表里有null就完蛋
    graph = (
        sess.g(
            oid_type="int64",
            vid_type="uint64",
            directed=False,
            generate_eid=True,
        )
        .add_vertices(
            acbus[["vid", "name"]],
            label="acbus",
            properties=None,
            vid_field="vid",
        )
        .add_vertices(
            acline[["vid", "name"]],
            label="acline",
            properties=None,
            vid_field="vid",
        )
        .add_vertices(
            acload[["vid", "name"]],
            label="acload",
            properties=None,
            vid_field="vid",
        )
        .add_vertices(
            acsource[["vid", "name"]],
            label="upper_grid",
            properties=None,
            vid_field="vid",
        )
        .add_vertices(
            storage[["vid", "name"]], label="storage", properties=None, vid_field="vid"
        )
    )

    bus_vertices_tbl = pd.concat(
        [acbus["vid"], acbus["name"]],
        axis=1,
        copy=False,
    )
    b1 = bus_vertices_tbl.rename(columns={"name": "bus1", "vid": "bus1vid"}, copy=False)
    b2 = bus_vertices_tbl.rename(columns={"name": "bus2", "vid": "bus2vid"}, copy=False)
    acline_vertices_tbl = pd.concat(
        [acline["vid"], acline["name"], acline["bus1"], acline["bus2"]],
        axis=1,
        copy=False,
    )
    df = pd.merge(acline_vertices_tbl, b1, how="inner", on="bus1", copy=False)
    df = pd.merge(df, b2, how="inner", on="bus2", copy=False)
    # graph不是inplace的
    graph = graph.add_edges(
        df[["vid", "bus1vid"]],
        label="connect",
        src_label="acline",
        dst_label="acbus",
        src_field="vid",
        dst_field="bus1vid",
        properties=[],
    ).add_edges(
        df[["vid", "bus2vid"]],
        label="connect",
        src_label="acline",
        dst_label="acbus",
        src_field="vid",
        dst_field="bus2vid",
        properties=[],
    )
    df = pd.merge(acload[["bus1", "vid"]], b1, how="inner", on="bus1", copy=False)
    graph = graph.add_edges(
        df[["vid", "bus1vid"]],
        label="connect",
        src_label="acload",
        dst_label="acbus",
        src_field="vid",
        dst_field="bus1vid",
        properties=[],
    )

    df = pd.merge(acsource[["bus1", "vid"]], b1, how="inner", on="bus1", copy=False)
    graph = graph.add_edges(
        df[["vid", "bus1vid"]],
        label="connect",
        src_label="upper_grid",
        dst_label="acbus",
        src_field="vid",
        dst_field="bus1vid",
        properties=[],
    )

    df = pd.merge(storage[["bus1", "vid"]], b1, how="inner", on="bus1", copy=False)
    graph = graph.add_edges(
        df[["vid", "bus1vid"]],
        label="connect",
        src_label="storage",
        dst_label="acbus",
        src_field="vid",
        dst_field="bus1vid",
        properties=[],
    )
    return sess, graph


def runIEEE4node(context: dict):
    # 借助graphscope建图，建完以后再使用Gremlin读取相邻顶点构造builder
    # 这样也不用自己merge来merge去的

    context["logger"].info("建立模型")
    sess, graph = build_grid(bus_table, acline_tbl, load_tbl, upper_table, storage_tbl)
    context["logger"].debug(f"grid schema:\n{graph.schema}")
    interactive = sess.gremlin(graph)
    node_num = interactive.execute("g.V().hasLabel('acbus').count()").one()
    context["logger"].debug(f"查询到{node_num}条交流母线")
    # res = interactive.execute("g.V().hasLabel('acbus').values('vid')").one()
    edge_num = interactive.execute("g.E().hasLabel('connect').count()").one()
    context["logger"].debug(f"查询到{edge_num}个连接边")

    model = Pregel(3)
    # 这狗屎查询就这样吧，受不了gremlin语言
    vids = interactive.execute("g.V().hasLabel('acline').values('vid')").one()
    out_v = [
        (
            interactive.execute(
                f"g.V().hasLabel('acline').has('vid',{id}).values('name')"
            ).one(),
            interactive.execute(
                f"g.V().hasLabel('acline').has('vid',{id}).out().values('vid')"
            ).one(),
        )
        for id in vids
    ]
    for k in range(0, len(vids)):
        model.set_Vertex(
            ADMM.acline_builder(
                vid=vids[k],
                name=out_v[k][0][0],
                out_vertices_id=out_v[k][1],
            )
        )

    vids = interactive.execute("g.V().hasLabel('acbus').values('vid')").one()
    out_v = [
        (
            interactive.execute(
                f"g.V().hasLabel('acbus').has('vid',{id}).values('name')"
            ).one(),
            interactive.execute(
                f"g.V().hasLabel('acbus').has('vid',{id}).out().values('vid')"
            ).one(),
        )
        for id in vids
    ]
    for k in range(0, len(vids)):
        model.set_Vertex(
            ADMM.bus_builder(
                vid=vids[k],
                name=out_v[k][0][0],
                out_vertices_id=out_v[k][1],
            )
        )

    vids = interactive.execute("g.V().hasLabel('upper_grid').values('vid')").one()
    out_v = [
        (
            interactive.execute(
                f"g.V().hasLabel('upper_grid').has('vid',{id}).values('name')"
            ).one(),
            interactive.execute(f"g.V().hasLabel('upper_grid').has('vid',{id}).out().values('vid')").one(),
        )
        for id in vids
    ]
    for k in range(0, len(vids)):
        model.set_Vertex(
            ADMM.acsource_builder(
                vid=vids[k],
                name=out_v[k][0][0],
                out_vertices_id=out_v[k][1],
            )
        )

    vids = interactive.execute("g.V().hasLabel('storage').values('vid')").one()
    out_v = [
        (
            interactive.execute(
                f"g.V().hasLabel('storage').has('vid',{id}).values('name')"
            ).one(),
            interactive.execute(
                f"g.V().hasLabel('storage').has('vid',{id}).out().values('vid')"
            ).one(),
        )
        for id in vids
    ]
    for k in range(0, len(vids)):
        model.set_Vertex(
            ADMM.storage_builder(
                vid=vids[k],
                name=out_v[k][0][0],
                out_vertices_id=out_v[k][1],
            )
        )

    vids = interactive.execute("g.V().hasLabel('acload').values('vid')").one()
    out_v = [
        (
            interactive.execute(
                f"g.V().hasLabel('acload').has('vid',{id}).values('name')"
            ).one(),
            interactive.execute(
                f"g.V().hasLabel('acload').has('vid',{id}).out().values('vid')"
            ).one(),
        )
        for id in vids
    ]
    for k in range(0, len(vids)):
        model.set_Vertex(
            ADMM.load_builder(
                vid=vids[k],
                name=out_v[k][0][0],
                out_vertices_id=out_v[k][1],
            )
        )

    context["logger"].info("建立模型成功")
    context["logger"].info("运行算例")
    model.run(context,200)


def start(work_dir: str):
    if not os.path.isdir(work_dir):
        raise ValueError("work_dir must be a directory")
    work_path = {"cwd": work_dir, "logd": os.path.join(work_dir, "log/")}
    if not os.path.exists(work_path["logd"]):
        os.makedirs(work_path["logd"])
    work_path.update(
        {
            "running_log": os.path.join(work_path["logd"], "running.log"),
            "iteration_logd": os.path.join(work_path["logd"], "iteration/"),
        }
    )
    logging.basicConfig(
        force=True, handlers=[logging.NullHandler()]
    )  # 取消输出到控制台
    logfile = logging.FileHandler(work_path["running_log"])
    console_log = logging.StreamHandler()
    logfile.setLevel(logging.DEBUG)
    logfile.setFormatter(
        logging.Formatter("%(asctime)s - %(name)s - %(levelname)s:\n%(message)s\n")
    )
    logger = logging.getLogger("admmLog")
    logger.addHandler(logfile)
    logger.removeHandler(console_log)
    logger.setLevel(logging.DEBUG)
    logger.info("start log")
    context = ADMM.context(
        rou=0.005,
        T=36,
        interval_seconds=600,
        tolerance=0.004,
        acline=acline_args_pd(line_table=acline_tbl, bus_table=bus_table),
        acload=acload_args_pd(load_tbl),
        storage=storage_args_pd(storage_tbl),
        upgrid=upper_args_pd(upper_table, 36),
        debug=True,
        logger=logger,
    )
    context["acline"].set_initialize_value(
        acline_args_pd.generate_zeros_initialize(context["acline"], context["T"])
    )
    context["storage"].set_initialize_value(
        storage_args_pd.generate_zeros_initialize(context["storage"], context["T"])
    )
    context["upgrid"].set_initialize_value(
        upper_args_pd.generate_zeros_initialize(context["upgrid"], context["T"])
    )
    runIEEE4node(context)

    logger.removeHandler(logfile)
    logfile.flush()
    logfile.close()


def is_update_converge(
    z_adj_terminals: ADMM.terminal_args, x_adj_terminals: ADMM.terminal_args, esp
) -> list[bool]:
    converged = [True for k in range(0, len(x_adj_terminals))]
    for k in range(0, len(x_adj_terminals)):
        if x_adj_terminals[k].var_status("P") == x_adj_terminals[k].args_status.INUSE:
            diff = np.abs(z_adj_terminals[k].P - x_adj_terminals[k].P)
            if not np.max(diff) <= esp:
                converged[k] = False
                continue

        if x_adj_terminals[k].var_status("Q") == x_adj_terminals[k].args_status.INUSE:
            diff = np.abs(z_adj_terminals[k].Q - x_adj_terminals[k].Q)
            if not np.max(diff) <= esp:
                converged[k] = False
                continue

        if x_adj_terminals[k].var_status("V") == x_adj_terminals[k].args_status.INUSE:
            diff = np.abs(z_adj_terminals[k].V2 - x_adj_terminals[k].V2)
            if not np.max(diff) <= esp:
                converged[k] = False
                continue
    return converged


def acline_update_native(
    context,
    name,
    terminal_messages: dict,
    history_messages: dict,
    bus1,
    bus2,
    iteration_data_for_acline: dict,
):

    value = ADMM.Vertex_acline.value_type(name)
    property_m = context["acline"].query(name)
    # 有新的用新的，没新的用旧的
    skip1 = True
    skip2 = True
    if terminal_messages[name].__contains__(bus1):
        terminal1 = terminal_messages[name][bus1]
        skip1 = False
    else:
        terminal1 = history_messages[name][bus1]

    if terminal_messages[name].__contains__(bus2):
        terminal2 = terminal_messages[name][bus2]
        skip2 = False
    else:
        terminal2 = history_messages[name][bus2]

    if skip1 is True and skip2 is True:
        # 没收到消息
        iteration_data_for_acline.update({name: None})
        return

    opt_sol = ADMM_update.acline_update(
        property_m["R"],
        property_m["X"],
        property_m["Vmax_i"],
        property_m["Vmin_j"],
        property_m["Imax"],
        terminal1,
        terminal2,
        context["rou"],
        context["T"],
    )

    # 更新是不分三七二十一都会把另一端叫起来
    terminal1.P = opt_sol["Pi"]
    terminal1.Q = opt_sol["Qi"]
    terminal1.V2 = opt_sol["V2i"]
    terminal1.LP = opt_sol["LPi"]
    terminal1.LQ = opt_sol["LQi"]
    terminal1.LV2 = opt_sol["LV2i"]
    terminal_messages[bus1][name] = terminal1
    history_messages[name][bus1] = terminal1

    terminal2.P = opt_sol["Pj"]
    terminal2.Q = opt_sol["Qj"]
    terminal2.V2 = opt_sol["V2j"]
    terminal2.LP = opt_sol["LPj"]
    terminal2.LQ = opt_sol["LQj"]
    terminal2.LV2 = opt_sol["LV2j"]
    terminal_messages[bus2][name] = terminal2
    history_messages[name][bus2] = terminal2

    terminal_messages[name].clear()
    for v in opt_sol.values():
        if isinstance(v, np.ndarray):
            v = v.reshape(-1).tolist()
    iteration_data_for_acline.update({name: opt_sol})
    return value


"""
写这个的时候发现建立的是二分图。
这样就形成了一个二步并行优化结构，在节点更新处判断收敛
"""


def node_update_native(
    context,
    name,
    terminal_messages: dict,
    history_messages: dict,
    sq,
    logger,
    iteration_data_for_acbus: dict,
):
    value = ADMM.Vertex_bus.value_type(name)
    if len(terminal_messages[name]) == 0:
        iteration_data_for_acbus.update({name: None})
        return

    # 合并端口变量
    adj_terminals = []
    for k in range(0, len(sq)):
        if terminal_messages[name].__contains__(sq[k]):
            adj_terminals.append(terminal_messages[name][sq[k]])
        else:
            adj_terminals.append(history_messages[name][sq[k]])

    up_adj_terminals = ADMM_update.node_update(
        adj_terminals, context["rou"], context["T"]
    )

    # 一致性收敛判断
    consensus_list = is_update_converge(
        up_adj_terminals,
        adj_terminals,
        context["tolerance"] * np.sqrt(context["T"] * len(sq)),
    )

    # 步进收敛判断
    if len(history_messages[name]) == len(sq):
        step_list = is_update_converge(
            up_adj_terminals,
            [history_messages[name][sq[k]] for k in range(0, len(sq))],
            context["tolerance"] * np.sqrt(context["T"] * len(sq)),
        )
    else:
        step_list = [False for k in range(0, len(sq))]

    for k in range(0, len(sq)):
        if consensus_list[k] is False or step_list[k] is False:
            terminal_messages[sq[k]][name] = up_adj_terminals[k]
            history_messages[name][sq[k]] = up_adj_terminals[k]
    terminal_messages[name].clear()

    info = dict(
        [
            (sq[k], up_adj_terminals[k].todict())
            for k in range(0, len(sq))
            if consensus_list[k] is False or step_list[k] is False
        ]
    )
    iteration_data_for_acbus.update({name: info})

    return value, up_adj_terminals, adj_terminals


def acline_init(context, name, terminal_messages, bus1, bus2):
    init_terminal = ADMM.terminal_args(context["T"], ("P", "Q", "V"))
    property_m = context["acline"].query(name)
    init_terminal.V2 = (
        np.ones((context["T"], 1), np.float64) * property_m["kV_bus1"] ** 2
    )
    terminal_messages[bus1].update({name: init_terminal})
    init_terminal.V2 = (
        np.ones((context["T"], 1), np.float64) * property_m["kV_bus2"] ** 2
    )
    terminal_messages[bus2].update({name: init_terminal})


def init_native(context: dict):
    terminal_messages = {
        "L12": {},
        "L23": {},
        "L34": {},
        "src": {},
        "sto1": {},
        "LD1": {},
        "b1": {},
        "b2": {},
        "b3": {},
        "b4": {},
    }

    acline_init(context, "L12", terminal_messages, "b1", "b2")
    acline_init(context, "L23", terminal_messages, "b2", "b3")
    acline_init(context, "L34", terminal_messages, "b3", "b4")

    init_terminal = ADMM.terminal_args(context["T"], ("P", "Q", "V"), ("V"))
    property_m = context["upgrid"].query("src")
    init_terminal.V2 = np.ones((context["T"], 1), np.float64) * (property_m["kV"] ** 2)
    terminal_messages["b1"].update({"src": init_terminal})

    init_terminal = ADMM.terminal_args(context["T"], ("P", "Q"))
    terminal_messages["b3"].update({"sto1": init_terminal})

    property_m = context["acload"].query("LD1")
    init_terminal = ADMM_update.load_update(
        property_m["Pt"],
        property_m["Qt"],
        ADMM.terminal_args(context["T"], ("P", "Q")),
        context["rou"],
        context["T"],
    )
    terminal_messages["b4"].update({"LD1": init_terminal})
    return terminal_messages


class AdmmJSONEncoder(json.JSONEncoder):
    def default(self, o):
        if isinstance(o, np.ndarray):
            return o.reshape(-1).tolist()
        return super().default(o)


def start_native(work_dir: str):
    if not os.path.isdir(work_dir):
        raise ValueError("work_dir must be a directory")
    work_path = {"cwd": work_dir, "logd": os.path.join(work_dir, "log/")}
    if not os.path.exists(work_path["logd"]):
        os.makedirs(work_path["logd"])
    work_path.update(
        {
            "running_log": os.path.join(work_path["logd"], "running.log"),
            "iteration_logd": os.path.join(work_path["logd"], "iteration/"),
        }
    )
    # 串行迭代，用来对比结果
    logging.basicConfig(
        force=True, handlers=[logging.NullHandler()]
    )  # 取消输出到控制台
    logfile = logging.FileHandler(work_path["running_log"])
    logfile.setLevel(logging.DEBUG)
    logfile.setFormatter(
        logging.Formatter("%(asctime)s - %(name)s - %(levelname)s:\n%(message)s\n")
    )
    logger = logging.getLogger("admmLog")
    logger.addHandler(logfile)
    logger.setLevel(logging.DEBUG)
    logger.info("start solving")

    context = ADMM.context(
        rou=0.005,
        T=36,
        interval_seconds=600,
        tolerance=0.004,
        acline=acline_args_pd(line_table=acline_tbl, bus_table=bus_table),
        acload=acload_args_pd(load_tbl),
        storage=storage_args_pd(storage_tbl),
        upgrid=upper_args_pd(upper_table, 36),
        debug=True,
        logger=logger,
    )
    max_iteration = 150
    terminal_messages = init_native(context)
    history_messages = {
        "L12": {},
        "L23": {},
        "L34": {},
        "src": {},
        "sto1": {},
        "LD1": {},
        "b1": {},
        "b2": {},
        "b3": {},
        "b4": {},
    }

    result_log = {}
    iter = 0
    converged = False
    while iter < max_iteration and converged is False:
        iter += 1
        # 求解母线优化问题，原始残差判断收敛,记录日志
        logger.debug(f"start superstep {iter}")
        result_log[iter] = {"acbus": {}, "acline": {}, "storage": {}, "up_grid": {}}
        node_update_native(
            context,
            "b1",
            terminal_messages,
            history_messages,
            ["src", "L12"],
            logger,
            result_log[iter]["acbus"],
        )

        node_update_native(
            context,
            "b2",
            terminal_messages,
            history_messages,
            ["L12", "L23"],
            logger,
            result_log[iter]["acbus"],
        )

        node_update_native(
            context,
            "b3",
            terminal_messages,
            history_messages,
            ["L23", "L34", "sto1"],
            logger,
            result_log[iter]["acbus"],
        )

        node_update_native(
            context,
            "b4",
            terminal_messages,
            history_messages,
            ["L34", "LD1"],
            logger,
            result_log[iter]["acbus"],
        )

        # 求解线路优化问题，记录日志
        acline_update_native(
            context,
            "L12",
            terminal_messages,
            history_messages,
            "b1",
            "b2",
            result_log[iter]["acline"],
        )
        acline_update_native(
            context,
            "L23",
            terminal_messages,
            history_messages,
            "b2",
            "b3",
            result_log[iter]["acline"],
        )
        acline_update_native(
            context,
            "L34",
            terminal_messages,
            history_messages,
            "b3",
            "b4",
            result_log[iter]["acline"],
        )
        # 求解储能优化问题，记录日志
        if len(terminal_messages["sto1"]) > 0:
            property_map = context["storage"].query("sto1")
            terminal1 = terminal_messages["sto1"]["b3"]
            y = ADMM_update.storage_update(
                PEssMax=property_map["Pmax"],
                eff_ch=property_map["eff_ch"],
                eff_dis=property_map["eff_dis"],
                capacity=property_map["capacity"],
                max_changing_N=property_map["max changing"],
                cost_operating=property_map["operating cost"],
                depreciation_per_change=property_map["changing cost"],
                interval_seconds=context["interval_seconds"],
                terminal=terminal1,
                rou=context["rou"],
                T=context["T"],
            )
            terminal1.P = y["Pi"]
            terminal1.Q = y["Qi"]
            terminal1.LP = y["LPi"]
            terminal1.LQ = y["LQi"]
            terminal_messages["b3"]["sto1"] = terminal1
            terminal_messages["sto1"].clear()
            result_log[iter]["storage"].update({"sto1": y})
        else:
            result_log[iter]["storage"].update({"sto1": None})

        # 给出负荷更新，记录日志
        if len(terminal_messages["LD1"]) > 0:
            property_map = context["acload"].query("LD1")
            update_res = ADMM_update.load_update(
                property_map["Pt"],
                property_map["Qt"],
                terminal_messages["LD1"]["b4"],
                context["rou"],
                context["T"],
            )
            terminal_messages["b4"]["LD1"] = update_res
            terminal_messages["LD1"].clear()
        # 给出上级电网更新，记录日志
        if len(terminal_messages["src"]) > 0:
            property_map = context["upgrid"].query("src")
            terminal1 = terminal_messages["src"]["b1"]
            y = ADMM_update.upper_grid_update(
                property_map["Ct"],
                property_map["Pmax"],
                property_map["Pmin"],
                terminal_messages["src"]["b1"],
                context["rou"],
                context["T"],
            )
            terminal1.P = y["Pi"]
            terminal1.Q = y["Qi"]
            terminal1.LP = y["LPi"]
            terminal1.LQ = y["LQi"]
            terminal1.V2 = np.ones((context["T"], 1), np.float64) * (
                property_map["kV"] ** 2
            )
            terminal_messages["b1"]["src"] = terminal1
            terminal_messages["src"].clear()
            result_log[iter]["up_grid"].update({"src": y})
        else:
            result_log[iter]["up_grid"].update({"src": None})

        # 判断收敛
        converged = True
        for k, v in terminal_messages.items():
            if len(v) > 0:
                converged = False
                break

    if not os.path.exists(work_path["iteration_logd"]):
        os.makedirs(work_path["iteration_logd"])
    result_log["summary"] = {"optimal value": None, "iteration": iter}

    for k in range(1, iter + 1):
        with open(
            os.path.join(work_path["iteration_logd"], f"iteration{k}.json"),
            "w",
            encoding="utf8",
        ) as f:
            json.dump(
                result_log[k], f, indent=4, ensure_ascii=False, cls=AdmmJSONEncoder
            )
    logger.info("finish")
    logger.removeHandler(logfile)
    logfile.flush()
    logfile.close()
